import os
import cv2
import numpy as np
import random
from tqdm import tqdm
import re
from glob import glob

def load_image(file_path):
    """Load a single image from the given path."""
    return cv2.imread(file_path, cv2.IMREAD_UNCHANGED)

def rotate_and_resize(image, mask, vis, angle, scale):
    """Rotate and resize an image, its mask, and its visualized mask."""
    h, w = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, scale)

    # Compute the bounding box for the rotated image
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    new_w = int((h * sin) + (w * cos))
    new_h = int((h * cos) + (w * sin))

    # Adjust the rotation matrix to account for the new dimensions
    M[0, 2] += (new_w / 2) - center[0]
    M[1, 2] += (new_h / 2) - center[1]

    # Perform the rotation and resizing
    rotated_image = cv2.warpAffine(image, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))
    rotated_mask = cv2.warpAffine(mask, M, (new_w, new_h), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    rotated_vis = cv2.warpAffine(vis, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))

    return rotated_image, rotated_mask, rotated_vis

def is_non_overlapping(mask, component_mask, x, y):
    """Check if a component can be placed at (x, y) without overlapping."""
    ch, cw = component_mask.shape[:2]
    if x + cw > mask.shape[1] or y + ch > mask.shape[0]:
        return False

    region = mask[y:y + ch, x:x + cw]
    region_non_zero = region > 0
    component_mask_non_zero = component_mask > 0
    overlap = np.sum(np.logical_and(region_non_zero, component_mask_non_zero))
    # overlap = np.sum(cv2.bitwise_and(region > 0, component_mask > 0))
    return overlap == 0

def place_component(board, board_mask, board_cpimask, board_vis_mask,
                    component, component_mask, cpmask, component_vis_mask,
                    x, y,feather=False):
    """Place a component on the board and update mask and visualized mask."""
    ch, cw = component.shape[:2]
    foreground_indices = component_mask > 0

     # Feather technique: Create a feathered mask
    if feather:
        feathered_mask = cv2.GaussianBlur(foreground_indices.astype(np.float32), (31, 31), 0)
        feathered_mask = feathered_mask / np.max(feathered_mask)
        for c in range(3):
            board[y:y + ch, x:x + cw, c][foreground_indices]=(
                feathered_mask[foreground_indices] * component[:, :, c][foreground_indices] +
                (1 - feathered_mask[foreground_indices]) * board[y:y + ch, x:x + cw, c][foreground_indices]
            ).astype(np.uint8)
        fb=board
    else:
        for c in range(3):
            board[y:y + ch, x:x + cw, c][foreground_indices] = component[:, :, c][foreground_indices]

    board_mask[y:y + ch, x:x + cw][foreground_indices] = np.maximum(
        board_mask[y:y + ch, x:x + cw][foreground_indices], component_mask[foreground_indices]
    )
    board_cpimask[y:y + ch, x:x + cw][foreground_indices] = np.maximum(
        board_cpimask[y:y + ch, x:x + cw][foreground_indices], cpmask[foreground_indices]
    )

    board_vis_mask[y:y + ch, x:x + cw][foreground_indices] = component_vis_mask[foreground_indices]

def synthesize_pcb_dataset(empty_board_path, component_root, output_folder, num_images, min_scale=0.5, max_scale=1.5,feather=True):
    """Generate synthetic PCB dataset."""
    def get_bbox_from_mask(mask, bgwh, shift):
        """
        从掩码图像中获取掩码部分的边界框。

        :param mask: 输入的掩码图像，形状为 (h, w)
        :return: 边界框信息 (x, y, width, height) xywh
        """
        # 找到掩码中非零像素的坐标
        rows, cols = np.nonzero(mask)
        bw, bh = bgwh
        sx, sy = shift

        if len(rows) == 0 or len(cols) == 0:
            return (0, 0, 0, 0)

        # 计算边界框的左上角坐标和宽高
        x = np.min(cols)
        y = np.min(rows)
        width = np.max(cols) - x + 1
        height = np.max(rows) - y + 1

        return ((x + width / 2 + sx) / bw, (y + height / 2 + sy) / bh, width / bw, height / bh)

    def get_component_type(component_path):
        type_name = component_path[len(component_root):]
        type_name = os.path.dirname(type_name)
        pattern = r'JRS(\d+)[a-zA-Z].*'
        match = re.match(pattern, type_name)
        number = match.group(1)
        return int(number)

    # Load the empty board
    pcb_bg_paths = glob(os.path.join(empty_board_path, "*.png"))
    bguseflag=[0] * len(pcb_bg_paths) 
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    component_paths = []  # len=1948
    for root, dirs, files in os.walk(component_root):
        for file in files:
            if file.endswith(('SolderLight.png', 'SolderLight.jpg', 'SolderLight.jpeg')) and 'label' not in root:
                component_paths.append(os.path.join(root, file))

    image_files = glob(os.path.join(output_folder, "image", "*.png"))
    start_index = len(image_files)
    for i in tqdm(range(start_index,num_images)):
        # pcb_bg_path = random.choice(pcb_bg_paths)

        # Choose pcb_bg_path based on usage flag
        min_use = min(bguseflag)
        min_indices = [index for index, value in enumerate(bguseflag) if value == min_use]
        chosen_index = random.choice(min_indices)
        pcb_bg_path = pcb_bg_paths[chosen_index]
        bguseflag[chosen_index] += 1  # Update the usage flag
        
        board = cv2.imread(pcb_bg_path)
        bh,bw,_=board.shape
        # max_n_cp = int(bh * bw / 300 / 400 * 1.2)
        # min_n_cp = int(bh * bw / 300 / 400 * 0.5)
        max_n_cp = int(bh * bw / 300 / 400 * 0.8)
        min_n_cp = int(bh * bw / 300 / 400 * 0.4)
        board_mask = np.zeros(board.shape[:2], dtype=np.uint8)
        board_cpimask = board_mask.copy()
        vis_mask = np.zeros_like(board)
        yolo_labels = []

        used_components = []
        num_components = random.randint(min_n_cp, max_n_cp)  # Random number of components per board
        error_guest = [
            "JRS183R0603/_a_SaveData__20241013161630397_Save Data [14837, 35647]",
            "JRS156R0603/_a_SaveData__20241013161630396_Save Data [22823, 35559]",
            "JRS157R0603/_a_SaveData__20241013161630397_Save Data [18766, 35613]",
            "JRS116PGLQFP/_a_SaveData__20241013155754673_SaveData [23447, 102811]",
            "JRS104TSOP50",
            "JRS116PGLQFP",
            "JRS82R1812/_a_SaveData__20241013155402990_Save Data [109592, 18847]",
        ]
        for it in range(num_components):
            component_path = random.choice(component_paths)
            label_path = component_path.replace('/_a_SaveData', '/label/_a_SaveData').replace('SolderLight.jpg', 'TrueLight.png')
            vis_path = component_path.replace('/_a_SaveData', '/label/_a_SaveData').replace('SolderLight.jpg', 'TrueLight_pseudo.png')

            log_file_path = 'og.log'
            error_flag = False
            for error_path in error_guest:
                if error_path in component_path:
                    it -= 1
                    error_flag = True
                    break
            if error_flag:
                continue

            component = load_image(component_path)
            component_mask = load_image(label_path)
            component_vis_mask = load_image(vis_path)

            if component is None or component_mask is None:
                it-=1
                warning_message = f"Warning: Failed to load data for {component_path}"
                status_message = f"{component is None}, {component_mask is None}"

                with open(log_file_path, 'a') as log_file:
                    log_file.write(warning_message + '\n')
                    log_file.write(status_message + '\n')
                continue

            angle = random.uniform(0, 360)
            scale = random.uniform(min_scale, max_scale)

            component, component_mask, component_vis_mask = rotate_and_resize(
                component, component_mask, component_vis_mask, angle, scale / 2
            )
            if feather:
                # Save feathered component
                feathered_board = board.copy()
                feathered_board_mask = board_mask.copy()
                feathered_board_cpimask = board_cpimask.copy()
                feathered_vis_mask = vis_mask.copy()

            ch, cw = component.shape[:2]
            do_it_more = max(1, 1024 * 1024 / ch / cw - 1)
            for _ in range(100):  # Maximum attempts to place the component
                x = random.randint(0, board.shape[1] - cw)
                y = random.randint(0, board.shape[0] - ch)

                if is_non_overlapping(board_mask, component_mask, x, y):
                    bbox_xywh = get_bbox_from_mask(component_mask, (bw,bh), (x, y))
                    class_id = get_component_type(component_path)

                    cpmask = np.copy(component_mask)
                    non_zero_indices = cpmask != 0
                    cpmask[non_zero_indices] = class_id

                    place_component(
                        board, board_mask, board_cpimask, vis_mask,
                        component, component_mask, cpmask, component_vis_mask,
                        x, y
                    )
                    # place_component(
                    #     feathered_board, feathered_board_mask, feathered_board_cpimask, feathered_vis_mask,
                    #     component, component_mask, cpmask, component_vis_mask,
                    #     x, y,feather=True
                    # )
                    used_components.append(component_path)
                    yolo_labels.append((class_id,) + bbox_xywh)
                    break

            del component, component_mask, component_vis_mask

        # Save outputs
        dir_list = ["image", "mask", "vis", "cplist", "cpbbox", "cpimask"]
        for dir_name in dir_list:
            if not os.path.exists(os.path.join(output_folder, dir_name)):
                os.makedirs(os.path.join(output_folder, dir_name))
        output_image_path = os.path.join(os.path.join(output_folder, "image"), f"pcb_{i:04d}.png")
        output_mask_path = os.path.join(os.path.join(output_folder, "mask"), f"pcb_{i:04d}_mask.png")
        output_cpimask_path = os.path.join(os.path.join(output_folder, "cpimask"), f"pcb_{i:04d}_cpimask.png")
        output_vis_mask_path = os.path.join(os.path.join(output_folder, "vis"), f"pcb_{i:04d}_vis_mask.png")
        output_components_path = os.path.join(os.path.join(output_folder, "cplist"), f"pcb_{i:04d}_components.txt")
        output_bbox_path = os.path.join(os.path.join(output_folder, "cpbbox"), f"pcb_{i:04d}_cpbbox.txt")
        cv2.imwrite(output_image_path, board)
        cv2.imwrite(output_mask_path, board_mask)
        cv2.imwrite(output_cpimask_path, board_cpimask)
        cv2.imwrite(output_vis_mask_path, vis_mask)
        
        # Save outputs for feathered board
        # output_feathered_image_path = os.path.join(os.path.join(output_folder, "image"), f"pcb_{i:04d}_feathered.png")
        # output_feathered_mask_path = os.path.join(os.path.join(output_folder, "mask"), f"pcb_{i:04d}_feathered_mask.png")
        # output_feathered_cpimask_path = os.path.join(os.path.join(output_folder, "cpimask"), f"pcb_{i:04d}_feathered_cpimask.png")
        # output_feathered_vis_mask_path = os.path.join(os.path.join(output_folder, "vis"), f"pcb_{i:04d}_feathered_vis_mask.png")
        # cv2.imwrite(output_feathered_image_path, feathered_board)
        # cv2.imwrite(output_feathered_mask_path, feathered_board_mask)
        # cv2.imwrite(output_feathered_cpimask_path, feathered_board_cpimask)
        # cv2.imwrite(output_feathered_vis_mask_path, feathered_vis_mask)


        with open(output_bbox_path, 'w') as f2:
            for label in yolo_labels:
                bbox_label = label[:5]
                label_str = " ".join(str(i) for i in bbox_label)
                f2.write(label_str + "\n")

        with open(output_components_path, 'w') as f:
            f.write('\n'.join(used_components))

# Example usage
# empty_board_path = "data_zoo/jrspcbbg"
# empty_board_path="data_zoo/simplepcb/train"
# component_root = "/home/lhx/pcbsyn/data_zoo/PCBA/20241209_PCBA样板标图数据-华科用/"
# output_folder = "synpcbmix"

empty_board_path='/home/lhx/pcbsyn/data_zoo/jrstruecolorbbg'
component_root='/home/lhx/pcbsyn/data_zoo/PCBA/20241209_PCBA样板标图数据含真彩-华科用/'
output_folder='syntruecolor'
synthesize_pcb_dataset(
    empty_board_path=empty_board_path,
    component_root=component_root,
    output_folder=output_folder,
    num_images=20,
    min_scale=1.1,
    max_scale=0.9
)